import re
import os
from collections import defaultdict
import fetch_detail
import reproduce
import json
import pandas as pd
import sys
import subprocess

from tools.logger import logger

# 项目结构
'''
./                              根目录
./oss-fuzz                      oss-fuzz源码


./project                       特定项目目录
./project/testcases             项目所有testcase
./project/project_source        项目github源码
./project/data                  所有打包数据
./project/work-n                项目复现测试目录

'''

task_build_path="/home/ranma/DataGenerator/task_build.xlsx"

def save_task_build_info(project_name, commit, fuzz_commit, bug,result):
    # 在原来的task_build_path文件中修改
    # 读取文件内容
    if os.path.exists(task_build_path):
        df = pd.read_excel(task_build_path)
    else:
        df = pd.DataFrame(columns=["commit", "project_name", "bug",  "fuzz_commit", "build", "reproduce"])
    # 添加一行
    # 根据commit bug的值查找是否已有存在
    if df[(df["commit"] == commit) & (df["bug"] == bug)].empty:
        new_row = pd.DataFrame({
            "commit": [commit],
            "project_name": [project_name],
            "fuzz_commit": [fuzz_commit],
            "bug": bug,
            "build": result["build"],
            "reproduce": result["reproduce"]
        })
        df = pd.concat([df, new_row], ignore_index=True)
    else:
        # 如果存在，则更新
        df.loc[(df["commit"] == commit) & (df["bug"] == bug), "build"] = result["build"]
        df.loc[(df["commit"] == commit) & (df["bug"] == bug), "reproduce"] = result["reproduce"]
    # 保存文件
    df.to_excel(task_build_path, index=False)



# 添加命令参数，用两个参数分别指定项目名称project_name以及项目git地址git_url
# 例如：python generate_task.py libxml2 https://github.com/GNOME/libxml2.git

# 读取命令行参数

# if len(sys.argv) != 3:
#     print("Usage: python generate_task.py project_name git_url")
#     exit(1)

# project_name = ""#sys.argv[1]
# git_url = sys.argv[2]
# flag = sys.argv[3]

# git_url="https://github.com/GNOME/libxml2.git"
# project_name = "libxml2"


def get_parent_commit(commit_hash, project_name):
    """
    获取指定提交的父提交。
    :param commit_hash: 提交的哈希值
    :return: 父提交的哈希值
    """
    try:
        # 使用git log命令获取父提交
        check_command=f'''
        cd ./{project_name}/{project_name} && git show --pretty=format:%H {commit_hash}^

        '''
        result = reproduce.run_command(check_command)
        
        parent_commit = result.strip().split('\n')[0]
        return parent_commit
    except subprocess.CalledProcessError as e:
        logger.error(f"Error getting parent commit for {commit_hash}: {e}")
        return None

def askfor_info(project_name):
    """
    print the info of the project
    """
    project_dir = os.path.join(os.getcwd(),project_name)
    file_path=f'{project_dir}/gitlog.txt'
    replog_path=f'{project_dir}/reprolog.txt'
    project_source = f"{project_dir}/{project_name}"
    testcases_dir = f"{project_dir}/testcases"
    data_dir = f"{project_dir}/data"

    print(f"TestCases for project {project_name}: ")
    for file in os.listdir(testcases_dir):
        print(file)
    # bug_dir=f"./oss-fuzz-vulns/vulns/{project_name}"
    

def build_env(project_name, git_url):
    # build_env用于获取project和osv的信息，并搭建相关环境
    fuzz_dir="./oss-fuzz"
    oss_fuzz_url="https://github.com/google/oss-fuzz.git"
    project_dir = os.path.join(os.getcwd(),project_name)
    file_path=f'{project_dir}/gitlog.txt'
    project_source = f"{project_dir}/{project_name}"
    testcases_dir = f"{project_dir}/testcases"
    data_dir = f"{project_dir}/data"
    bug_dir=f"./oss-fuzz-vulns/vulns/{project_name}"

    pattern_for_commit = re.compile(r'[a-f0-9]{40}')
    pattern_for_introduced = re.compile(r'introduced: [a-f0-9]{40}')
    pattern_for_fixed = re.compile(r'fixed: [a-f0-9]{40}')
    pattrn_for_url = re.compile(r"OSS-Fuzz report: (https?://[^\s]+)")


    # build the basic structure
    if not os.path.exists(project_dir):
        os.mkdir(project_dir)
    if not os.path.exists(testcases_dir):
        os.mkdir(testcases_dir)
    if not os.path.exists(data_dir):
        os.mkdir(data_dir)

    # clone the oss-fuzz repo
    if not os.path.exists(fuzz_dir):
        os.system(f"git clone {oss_fuzz_url} {fuzz_dir}")

    # clone the project source code
    if not os.path.exists(project_source):
        os.system(f"git clone {git_url} {project_source}")
    logger.debug(f"Finish cloning {project_name} and oss-fuzz")

    # 读取文件内容
    # 在项目源码下调用git log命令，并获取命令返回内容，将结果保存到gitlog.txt文件中
    content=os.popen(f"cd {project_source} && git log --format=%H ").read()
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(content)
    # 使用正则表达式匹配commit号
    commit_hashes = re.findall(pattern_for_commit, content)
    # reverse commit_hashes
    commit_hashes.reverse()
    # 截取commit号的后40位
    commit_hashes = [commit_hash[-40:] for commit_hash in commit_hashes]
    logger.debug(f"Finish fetching commit hashes for {project_name}")

    # 打印所有找到的commit号
    # for commit_hash in commit_hashes:
    #     print(commit_hash)

    bug_map=defaultdict(dict)

    if os.path.exists(f"{project_dir}/bug_{project_name}.json"):
        with open(f"{project_dir}/bug_{project_name}.json", 'r', encoding='utf-8') as file:
            bug_map=json.load(file)


    for root, dirs, files in os.walk(bug_dir):
        for file in files:
            file_path = os.path.join(root, file)

            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            introduced = re.findall(pattern_for_introduced, content)[0][-40:]
            fixed = re.findall(pattern_for_fixed, content)
            url = re.findall(pattrn_for_url, content)[0]

            if fixed!=[]:
                fixed=fixed[0][-40:]
            else :
                fixed=commit_hashes[-1]
            file_name = file_path.split("/")[-1].split(".")[0]   
            logger.info(f"file: {file_name} introduced: {introduced} fixed: {fixed}")  
            #print(f"file: {file_name} introduced: {introduced} fixed: {fixed}")
            # 遍历commit_hashes，打印从fixed到introduced之间的commit号


            # fetch the detail for reproduce
            # check if file_name in bug_map's keys
            if file_name not in bug_map: 
                try :            
                    detail=fetch_detail.get_details(issue_url=url)
                    logger.debug(f"Finish fetching details for {file_name}")
                except Exception as e:
                    logger.debug(f"Fail to fetch details for {file_name} {e}") 
                    continue
                bug_map[file_name] = {"introduced": introduced, "fixed": fixed, "url": url, "git": git_url}
                bug_map[file_name]["reproduce"]=detail
                with open(f"{project_dir}/bug_{project_name}.json", 'w', encoding='utf-8') as file:
                    json.dump(bug_map, file, ensure_ascii=False, indent=4)

            if not os.path.exists(f"{testcases_dir}/testcase-{file_name}"):
                with open(f"{testcases_dir}/testcase-{file_name}","wb") as f :
                    try:
                        f.write(fetch_detail.get_testcase(detail["reproducer_testcase"]))
                        logger.debug(f"Finish downloading testcase for {file_name}")
                    except Exception as e:
                        logger.debug(f"Fail to download testcase for {file_name} {e}") 


    # save bug_map to json file
    
    with open(f"{project_dir}/bug_{project_name}.json", 'w', encoding='utf-8') as file:
        json.dump(bug_map, file, ensure_ascii=False, indent=4)

    # with open(f"{project_dir}/bug_{project_name}.json", 'r', encoding='utf-8') as file:
    #     bug_map=json.load(file)

def generate_task(project_name, git_url, commit, fuzz_commit, latest, oldest, bugs):
    # 先写task_info.json文件
    # 创建一个以commit为名的文件夹
    # 将上述的几个tar.gz文件放入文件夹
    oss_fuzz_url="https://github.com/google/oss-fuzz.git"
    fuzz_dir=os.path.join(os.getcwd(),"oss-fuzz")
    project_dir = os.path.join(os.getcwd(),project_name)
    project_source = f"{project_dir}/{project_name}"
    testcases_dir = f"{project_dir}/testcases"
    data_dir = f"{project_dir}/data"
    bug_map_path = f"{project_name}/bug_{project_name}.json"
    with open(bug_map_path, 'r', encoding='utf-8') as file:
        bug_map = json.load(file)
    os.chdir(data_dir)
    workdir=f"{commit}"
    # 1. 创建文件夹
    os.mkdir(workdir)
    original_dir = os.getcwd()
    os.chdir(workdir)

    
    # os.system(f"git clone {git_url} ./{project_name}")
    # # 1. 创建 DIFF.tar.gz
    # # diff文件的内容为：commit号和bug对应的fixed号的diff
    # # 将commit和introduced的parent commit之间的diff保存为DIFF.diff

    # # fixed =bug_map[bug_name]["introduced"]

    # #fixed = bug_map[bug_name]["fixed"]
    # os.system(f"cd ./{project_name} && git diff --binary {oldest} {latest} > ../DIFF.diff")
    # os.system(f"tar -czf ./DIFF.tar.gz ./DIFF.diff")

    # # 2. 创建 REPO.tar.gz
    # os.system(f"cd ./{project_name} && git checkout {oldest}")
    # os.system(f"tar -czf ./{project_name}.tar.gz ./{project_name}")

    # # 3. 创建 fuzz-tooling.tar.gz
    # os.system(f"git clone {oss_fuzz_url} ./fuzz-tooling")
    # os.system(f"cd ./fuzz-tooling && git checkout {fuzz_commit}")
    # # 删去oss_fuzz下projects文件夹中除了libxml2之外的所有文件
    # os.system(f"cd ./fuzz-tooling/projects && find . -mindepth 1 -maxdepth 1 ! -name {project_name}"+" -exec rm -rf {} \;")

    # os.system(f"tar -czf ./fuzz-tooling.tar.gz ./fuzz-tooling")

    
    task_info=defaultdict(dict)
    task_info["bugs"]=[]
    # 4. 创建一个json文件，包含所有相关信息
    if os.path.exists(f"./task_info.json"):
        with open(f"./task_info.json", 'r', encoding='utf-8') as file:
            task_info = json.load(file)
    # 合并task_info["bugs"] 与 bugs，并去重
    set1 = set(task_info["bugs"])
    set2 = set(bugs)
    bugs = list(set1.union(set2))

    task_info = {
        "git_url": git_url,
        "commit": commit,
        "fuzz_commit": fuzz_commit,
        "latest": latest,
        "oldest": oldest,
        "bugs": bugs
    }
    bug_info_new = {}
    for bug in bugs:
        bug_info_new[bug] = bug_map[bug]

    with open(f"./task_info.json", 'w', encoding='utf-8') as file:
        json.dump(task_info, file, ensure_ascii=False, indent=4)
    with open('./bug_info.json', 'w') as f:
        json.dump(bug_info_new, f, indent=4)

    # 5. 将testcase复制到TEST目录下
    if not os.path.exists(f"./testcases"):
        os.mkdir(f"./testcases")
    for bug in bugs:
        os.system(f"cp {testcases_dir}/testcase-{bug} ./testcases")
    os.chdir(original_dir)

def reproduce_bug(project_name, bug):

    # related path
    fuzz_dir=os.path.join(os.getcwd(),"oss-fuzz")
    project_dir = os.path.join(os.getcwd(),project_name)
    project_source = f"{project_dir}/{project_name}"
    testcases_dir = f"{project_dir}/testcases"
    data_dir = f"{project_dir}/data"

    # load bug map
    bug_map_path = f"{project_name}/bug_{project_name}.json"
    with open(bug_map_path, 'r', encoding='utf-8') as file:
        bug_map = json.load(file)
    
    if bug not in bug_map:
        logger.warning(f"Bug {bug} not found")
        return

    # load message for reproduce
    fuzz_map=reproduce.parse_git_log(reproduce.get_git_log(fuzz_dir),1)
    project_map=reproduce.parse_git_log(reproduce.get_git_log(project_source),2)
    commit = bug_map[bug]["introduced"]
    fuzz_commit=reproduce.find_suite_commit(commit,fuzz_map,project_map)
    logger.info(f"commit: {commit}")
    logger.info(f"fuzz_commit: {fuzz_commit}")

    detail = bug_map[bug]["reproduce"]
    git_url = bug_map[bug]["git"]
    testcases_dir = f"{project_name}/testcases"
    work_dir = f"{project_name}/work-rep"
    if not os.path.exists(work_dir):
        os.mkdir(work_dir)
    result = reproduce.reproduce_once(commit, fuzz_commit, detail, bug, testcases_dir, work_dir)
    logger.debug(f"result: {result}")
    if result["build"]:
        logger.info(f"reproduce success for {bug}")
        latest = bug_map[bug]["introduced"]
        oldest = bug_map[bug]["introduced"]
        oldest = get_parent_commit(oldest, project_name)
        logger.info(f"latest: {latest} oldest: {oldest}")
        generate_task(project_name, git_url, commit, fuzz_commit, commit, oldest, [bug])
    else :
        logger.info(f"reproduce failed for {bug}")
    save_task_build_info(project_name, commit, fuzz_commit, bug, result)
    return result




# # x=input()
# # print("OK")

# # # 遍历所有commit号，打印每个commit号对应map的长度，并挑出最长的一个
# # max_len = 0
# # max_commit = ""
# # for commit in commit_bug_map:
# #     print(f"commit: {commit} len: {len(commit_bug_map[commit])}")
# #     if len(commit_bug_map[commit]) > max_len:
# #         max_len = len(commit_bug_map[commit])
# #         max_commit = commit

# # print(f"max_commit: {max_commit} max_len: {max_len}")

# # 将这个map保存为excel文件，每一行为commit，每一列为文件名，内容为1或0，并添加一列统计同和
# def save_map_excel(tosave_map, excel_name):
#     df = pd.DataFrame(tosave_map).T.fillna(0)
#     df['sum'] = df.sum(axis=1)
#     df.to_excel(f"{project_dir}/{excel_name}")
#     #print(df)    



# # 对于每个任务（可以包含一个或多个漏洞），准备：
# # - REPO.tar.gz （包含项目的源代码）
# # - OSS_FUZZ.tar.gz （包含 OSS-fuzz 仓库中用于构建项目的代码）
# # - DIFF.tar.gz （包含一个 diff 文件）
# # 创建一个函数用于进行上述工作
# # 函数参数包括：commit号，bug名称
# # 函数返回值包括：REPO.tar.gz，OSS_FUZZ.tar.gz，DIFF.tar.gz



    

# def choose_few_cases(number, data, ascending=False, exclude=[], exclude_commit=[]):
#     # 选取数量最多的一批commit
#     df = pd.DataFrame(data)
#     df = pd.DataFrame(data).T.fillna(0)
#     # 删掉exclude中的列名
#     df = df.drop(columns=exclude)

#     #删除exclude_commit中的行
#     df = df.drop(exclude_commit, axis=0)

#     if 'sum' not in df.columns:
#         df['sum'] = df.sum(axis=1)

#     # 根据'sum'列进行排序（降序）
#     sorted_df = df.sort_values(by='sum', ascending=ascending)

#     # 选择前十行
#     top_num = sorted_df.head(number)

#     # 将前十行转换为字典
#     top_num = top_num.drop(columns=['sum'])
#     top_num_dict = top_num.to_dict(orient='index')
#     print("hello_hello", top_num_dict)
#     return top_num_dict

   


# # 将漏洞复现流程自动化
# # 对于同一个commit
# # 的版本，通过commit_bug_map查看其存在的bug
# # 构建针对这个commit版本的复
# # 现测试环境
# # 将每一个bug的testcase进行运行，查看是否能跑出ERROR
# # 将结果进行保存
# work_dir=f"{project_dir}/work-rep"
# if not os.path.exists(work_dir):
#     os.mkdir(f"{project_dir}/work-rep")
# fuzz_map=reproduce.parse_git_log(reproduce.get_git_log(fuzz_dir),1)
# project_map=reproduce.parse_git_log(reproduce.get_git_log(project_source),2)
# commit_choose=choose_few_cases(8, commit_bug_map, False)
# save_map_excel(commit_choose, "overlap.xlsx")
# success_repro=[]
# failed_commit=[]
# temp_map=commit_bug_map.copy()

# # 第一轮，尽可能选bug多的case
# if flag=="single":
#     round = 0
# else:
#     round = 25
# while len(success_repro)<len(bug_map) and round>0:
#     round-=1
#     commit_choose=choose_few_cases(1, temp_map, False, success_repro, failed_commit)
#     if commit_choose==[]:
#         break
#     print(commit_choose)
#     # 检查f"{data_dir}/{commit_choose}"是否存在，如果存在则跳过
#     commit_name=list(commit_choose.keys())[0]
#     if os.path.exists(f"{data_dir}/{commit_name}"):
#         # 读取其中的task_info.json文件，获取bug列表并加入到success_repro中
#         print(f"{commit_name} already exists")
#         with open(f"{data_dir}/{commit_name}/task_info.json", 'r', encoding='utf-8') as file:
#             task_info = json.load(file)

#         success_repro.extend(task_info["bugs"])
#         continue

#     for commit in commit_choose:
#         fuzz_commit=reproduce.find_suite_commit(commit,fuzz_map,project_map)
#         bug_list=[]
#         for bug in commit_choose[commit]:
#             if commit_choose[commit][bug]==1:   
#                 # print(fuzz_commit,"time: ",fuzz_map[fuzz_commit])
#                 print(f"trying to reproduce bug {bug} on commit {commit}")
#                 if not reproduce.reproduce_once_d(commit,fuzz_commit,bug_map[bug]["reproduce"],bug,git_url,testcases_dir,work_dir):
#                     commit_choose[commit][bug]=-1
#                     print("reproduce failed!!!")
#                 else:
#                     commit_choose[commit][bug]=2
#                     bug_list.append(bug)
#                     success_repro.append(bug)
#                     print("reproduce success!!!")
#                 save_map_excel(commit_choose, "reproduce.xlsx")
        
#         # 得到了bug_list以后，选取其中所有bug的introdeced commit中最新和最老的
#         # 根据这些commit在commit_hashes中的位置，选取最靠前和最后的commit
#         # 初始commit是bug_list[0]的introduced commit
#         if bug_list!=[]:
#             latest = bug_map[bug_list[0]]["introduced"]
#             oldest = bug_map[bug_list[0]]["introduced"]
#             for bug in bug_list:
#                 if commit_hashes.index(bug_map[bug]["introduced"]) > commit_hashes.index(latest):
#                     latest = bug_map[bug]["introduced"]
#                 if commit_hashes.index(bug_map[bug]["introduced"]) < commit_hashes.index(oldest):
#                     oldest = bug_map[bug]["introduced"]
            
            
#             oldest = commit_hashes[commit_hashes.index(oldest) + 1]
#             print(f"latest: {latest} oldest: {oldest}")
#             generate_task(commit, fuzz_commit, commit, oldest, bug_list)
#         else:
#             failed_commit.append(commit)


# bug_third=defaultdict(dict)
# # 第二轮，逐个bug选取
# if flag=="single":
#     # 读取data_dir下的所有文件夹，将其中的task_info.json文件读取，获取其中的bug列表
#     for root, dirs, files in os.walk(data_dir):
#         for dir in dirs:
#             if os.path.exists(f"{data_dir}/{dir}/task_info.json"):
#                 with open(f"{data_dir}/{dir}/task_info.json", 'r', encoding='utf-8') as file:
#                     task_info = json.load(file)
#                 success_repro.extend(task_info["bugs"])
# print("first round finished, ", success_repro)



# # 第三轮，使用libFuzzer进行fuzzing
# print("Second round finished, ", success_repro)
# for bug in bug_third:
#     if bug not in success_repro :#and bug == "OSV-2020-1623":

#         print(f"trying to reproduce bug {bug}")
#         commit=bug_map[bug]["introduced"]
#         if os.path.exists(f"{data_dir}/{commit}"):
#             print(f"{commit} already exists")
#             with open(f"{data_dir}/{commit}/task_info.json", 'r', encoding='utf-8') as file:
#                 task_info = json.load(file)
#             success_repro.extend(task_info["bugs"])
#             continue
#         fuzz_commit=reproduce.find_suite_commit(commit,fuzz_map,project_map)
#         if not reproduce.reproduce_once_d(commit,fuzz_commit,bug_third[bug]["reproduce"],bug,git_url,testcases_dir,work_dir):
#             print("reproduce failed!!!")
#         else:
#             success_repro.append(bug)
#             print("reproduce success!!!")
#             latest = bug_map[bug]["introduced"]
#             oldest = bug_map[bug]["introduced"]
#             oldest = commit_hashes[commit_hashes.index(oldest) + 1]
#             generate_task(commit, fuzz_commit, commit, oldest, [bug])
#             save_map_excel(commit_choose, "reproduce.xlsx")

# with open(replog_path,"a") as f:
#     f.write(
#         f"current reproduce success: {success_repro}\n"    
#     )







